import argparse
import torch
from PIL import Image
from torch import nn,Tensor
from torchvision.models import resnet18
from torchvision.transforms import transforms
from torchvision.io import read_image

normalize = transforms.Normalize([0.485, 0.456, 0.406],
                                 [0.229, 0.224, 0.225])

transform_data = transforms.Compose([
            transforms.Resize((64, 64)),
            transforms.ToTensor(),
            normalize
        ])

Classes:list = [
    "America_Ferrera",
    "Christina_Applegate",
    "Colin_Firth",
    "Courteney_Cox",
    "Daniel_Day-Lewis",
    "Debra_Messing",
    "Emile_Hirsch",
    "Ethan_Hawke",
    "Felicity_Huffman",
    "Fran_Drescher",
    "Gabriel_Macht",
    "Geena_Davis",
    "Glenn_Close",
    "Holly_Marie_Combs",
    "Jon_Voight",
    "Julie_Benz",
    "Kristin_Chenoweth",
    "Matt_Damon",
    "Neve_Campbell",
    "Tina_Fey"
    ]

def load_model(model_path:str):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = resnet18(num_classes = 20).to(device)
    model.load_state_dict(torch.load(model_path))
    return model,device

def pil_loader(pth:str):
    with open(pth, 'rb') as f:
        img = Image.open(f)
        return img.convert('RGB')

def load_data(data_path) -> Tensor: 
    image = pil_loader(data_path)
    image_tensor:Tensor = transform_data(image)
    image_tensor = image_tensor.unsqueeze(0)
    return image_tensor

def eval_data(x:Tensor,model:nn.Module,device:str,classes:list):
    model.eval()
    with torch.no_grad():
        x = x.to(device)
        pred = model(x)
        predicted = classes[pred[0].argmax(0)]
        print(f'Predicted: "{predicted}"')

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="DPFL experiment")
    parser.add_argument('--model_path', type=str, default="results/central/cen/resnet18_facescrub.pth")
    parser.add_argument('--data_path', type=str, default="../data/facescrub/test/America_Ferrera/America_Ferrera_54569_25821.jpeg")
    args = parser.parse_args()

    model_pth = args.model_path
    data_pth = args.data_path

    model,device = load_model(model_pth)
    image_tensor:Tensor = load_data(data_pth)
    eval_data(image_tensor,model,device,Classes)